"""Script model to ONNX format"""

import torch

from src.initial_model_utils import init_model

MODEL_WEIGHTS = "./data/weights/resnet50_tl_68.pt"

device = torch.device("cpu")
model = init_model(device, num_classes=35)
model.load_state_dict(torch.load(MODEL_WEIGHTS, map_location=device))
model.eval()

# TODO
# Дописать конвертацию в ONNX
